from config import DeepSeekVanna
from config import DatabaseManager
from config import config
import mysql.connector
import streamlit as st
import os

# 初始化组件
db_manager = DatabaseManager()
vn = DeepSeekVanna(db_manager)

# 初始化数据库连接
vn.conn = mysql.connector.connect(
    host=config.db_host,
    port=config.db_port,
    user=config.db_user,
    password=config.db_password,
    database=config.db_name
)

# 训练模型
def train_model():
    # 注入DDL
    with db_manager.get_connection() as conn:
        schema = vn._get_schema(conn)
    vn.train(ddl=schema)
    
    # 注入示例SQL（从文件加载）
    # 处理sql_examples.txt（旧格式）
    with open("train_data/sql_examples.txt", encoding='utf-8-sig') as f:
        for line in f:
            parts = line.strip().split("|||")
            if len(parts) != 2:
                raise ValueError(f"训练数据格式错误: {line}")
            question, sql = parts
            vn.train(question=question.strip(), sql=sql.strip())
    
    # 处理text_examples.txt（新格式）
    with open("train_data/text_examples.txt", encoding='utf-8-sig') as f:
        current_question = None
        current_sql = None
        for line in f:
            line = line.strip()
            if line.startswith("用户问题:"):
                current_question = line.split(":", 1)[1].strip()
            elif line.startswith("SQL:"):
                current_sql = line.split(":", 1)[1].strip().rstrip(';')
                if current_question and current_sql:
                    vn.train(question=current_question, sql=current_sql)
                    current_question = None
                    current_sql = None

    # 处理背景文档（支持.txt和.md格式）
    doc_count = 0
    docs_path = "train_data/docs"
    if os.path.exists(docs_path):
        for filename in os.listdir(docs_path):
            if filename.endswith((".txt", ".md")):
                filepath = os.path.join(docs_path, filename)
                try:
                    with open(filepath, "r", encoding="utf-8") as doc_file:
                        doc_content = doc_file.read()
                        vn.train(documentation=doc_content)
                        doc_count += 1
                        st.toast(f"已训练文档: {filename}", icon="📄")
                except UnicodeDecodeError:
                    st.warning(f"文档编码错误: {filename}（请使用UTF-8编码）")
                except Exception as e:
                    st.warning(f"文档处理失败 [{filename}]: {str(e)}")
    
    st.toast(f"完成文档训练，共处理 {doc_count} 个背景文档", icon="✅")

# Streamlit界面
st.title("DeepSeek NL2SQL 助手")
if st.button("初始化训练"):
    with st.spinner("模型训练中..."):
        train_model()
    st.success("训练完成！")

question = st.text_input("输入业务问题，例如：'显示最近一个月销售额TOP10的产品'")
if question:
    try:
        # 流式生成响应
        response_stream = vn.generate_sql(question, stream=True)
        response_text = ""
        progress_bar = st.progress(0)
        status_text = st.empty()
        sql_block = st.empty()
        comment_display = st.empty()
        res_text = st.empty()
        final_sql = ""
        
        # 实时处理流式输出
        for chunk in response_stream:
            response_text += chunk
            progress = min(len(response_text) / 500, 1.0)
            
            # 更新进度条和状态
            progress_bar.progress(progress)
            status_text.info(f"生成进度: {int(progress*100)}%")
            
            # 流式解析SQL块
            import re
            sql_match = re.search(r'```sql\s*(.*?)\s*```', response_text, re.DOTALL)
            if sql_match:
                final_sql = sql_match.group(1).strip()
                sql_block.code(f"生成SQL：\n```sql\n{final_sql}\n```", language="sql")
                
            # 流式解析注释
            comment_match = re.search(r'^([^`]+?)(?=```sql)', response_text, re.DOTALL)
            if comment_match:
                comment = comment_match.group(1).strip()
                comment_display.info(f"分析说明：\n{comment}")
        res_text.info(f"Deepseek原始信息：{response_text}")
        # 隐藏进度组件
        progress_bar.empty()
        status_text.empty()
        
        # 最终校验SQL
        if not final_sql:
            st.error(f"未找到有效SQL代码块，原始响应内容：\n{response_text}")
            st.stop()
            
        # 清理SQL中的转义字符
        sql = re.sub(r'\\`', '`', final_sql)
        
        # SQL语法校验
        from sqlvalidator import format_sql, parse
        try:
            parsed_sql = format_sql(sql)
            if not parse(parsed_sql):
                st.error("SQL语法校验失败：\n" + "\n".join(parsed_sql.errors))
                st.stop()
        except Exception as e:
            st.warning(f"校验器异常：{str(e)}，继续执行原始SQL")
            
        cursor = None
        try:
            cursor = vn.conn.cursor(dictionary=True)
            cursor.execute(sql)
            df = cursor.fetchall()
        except Exception as e:
            raise RuntimeError(f"SQL执行错误: {str(e)}") from e
        finally:
            if cursor:
                cursor.close()
        st.dataframe(df)
        
        # 自动可视化（转换为DataFrame）
        if len(df) > 0:
            import pandas as pd
            df = pd.DataFrame(df)
            if len(df.columns) >= 1:
                # 使用columns布局优化显示
                col1, col2 = st.columns([3, 2])
                with col1:
                    st.dataframe(df)
                with col2:
                    try:
                        st.bar_chart(df.set_index(df.columns[0]))
                    except Exception as e:
                        st.warning(f"可视化失败: {str(e)}")
    except Exception as e:
        error_msg = f"SQL执行错误: {str(e)}"
        st.error(error_msg)
        
        # 将错误信息加入对话历史并重新生成
        if "chat_history" not in st.session_state:
            st.session_state.chat_history = []
        st.session_state.chat_history.append({
            "question": question,
            "sql": sql,
            "error": error_msg
        })
        
        with st.spinner("正在根据错误信息重新生成SQL..."):
            # 携带历史对话重新生成
            new_response = vn.generate_sql(
                question=question,
                chat_history=st.session_state.chat_history,
                stream=True
            )
            
            # 重新处理生成结果
            new_response_text = ""
            new_final_sql = ""
            
            # 复用原有的显示组件
            progress_bar = st.progress(0)
            status_text.info("重新生成中...")
            
            for chunk in new_response:
                new_response_text += chunk
                progress = min(len(new_response_text) / 500, 1.0)
                progress_bar.progress(progress)
                
                # 解析新SQL
                new_sql_match = re.search(r'```sql\s*(.*?)\s*```', new_response_text, re.DOTALL)
                if new_sql_match:
                    new_final_sql = new_sql_match.group(1).strip()
                    sql_block.code(f"重新生成SQL：\n```sql\n{new_final_sql}\n```", language="sql")
                    
                # 解析新注释
                new_comment_match = re.search(r'^([^`]+?)(?=```sql)', new_response_text, re.DOTALL)
                if new_comment_match:
                    new_comment = new_comment_match.group(1).strip()
                    comment_display.info(f"重新分析说明：\n{new_comment}")
            
            # 清理组件
            progress_bar.empty()
            status_text.empty()
            
            if new_final_sql:
                # 执行新SQL
                try:
                    cursor = vn.conn.cursor(dictionary=True)
                    cursor.execute(sql)
                    df = cursor.fetchall()
                    st.dataframe(df)
                    st.success("重新生成SQL执行成功！")
                except Exception as e:
                    st.error(f"重新生成SQL仍存在错误: {str(e)}")
                finally:
                    if cursor:
                        cursor.close()
